import datetime
from langchain_deepseek import ChatDeepSeek
from langchain.agents import create_agent
from langchain_core.tools import tool
from langgraph.checkpoint.memory import InMemorySaver

checkpointer = InMemorySaver()  # 创建检查点对象, 用于保存对话上下文


@tool
def get_current_date(input: str):
    """获取今天日期
    Returns:
        str: 今天的日期，格式为 YYYY-MM-DD
    """
    return datetime.datetime.today().strftime("%Y-%m-%d")


# 配置会话参数，使用固定的线程ID保持对话上下文
config = {  # 配置会话参数
    "configurable": {
        "thread_id": "1"  # 固定线程ID
    }
}

# 初始化DeepSeek模型客户端
llm = ChatDeepSeek(
    model="deepseek-chat",  # 指定模型名称
    api_key="sk-dc3f063534c542a3a403760655d3eae2",  # 替换为你的API Key
    temperature=0.7,  # 控制生成结果的随机性，0表示更确定性的输出
    max_tokens=None,  # 生成内容的最大长度，None代表使用模型默认值
    timeout=60,  # 请求超时设置
    max_retries=2,  # 请求失败重试次数
)
# 创建AI助手代理，配置模型、工具和系统提示词
agent = create_agent(
    model=llm,  # 使用 DeepSeek 模型
    tools=[get_current_date],  # 添加工具
    checkpointer=checkpointer,  # 使用检查点对象保存对话上下文
    system_prompt="你是一个乐于助人的助手。当用户询问当前日期时，请使用get_current_date工具。"  # 系统提示词
)
res1 = agent.invoke({"messages": [{"role": "user", "content": "使用Python写一个冒泡排序算法"}]}, config)
print("res1:", res1)

res2 = agent.invoke({"messages": [{"role": "user", "content": "java呢"}]}, config)
print("res2:", res2)
